#!/usr/bin/env python3
"""
Point Cloud Trimming Script

This script trims point cloud data based on the focused_bbox defined in pcdManip.json.
It filters out points that fall outside the specified bounding box.

Usage:
    This module provides trimPcd function for trimming single frame point cloud data
    within the unified data format based on bounding box constraints.
"""

import numpy as np
import zarr
import argparse
import os
import json
import logging
from pathlib import Path
from typing import Tuple, Optional, Union, List

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class BboxPointCloudTrimmer:
    """
    Trims point cloud data based on bounding box constraints.
    """
    
    def __init__(self, bbox_min: np.ndarray, bbox_max: np.ndarray):
        """
        Initialize the trimmer with configuration from pcdManip.json.
        
        Args:
            config_path: Path to pcdManip.json config file. If None, uses default location.
        """
        # logger.info(f"Bbox min: {bbox_min}")
        # logger.info(f"Bbox max: {bbox_max}")
        self.bbox_min = bbox_min
        self.bbox_max = bbox_max
    #     if config_path is None:
    #         config_path = "pcdManip.json"
        
    #     self.load_config()
        
    # def load_config(self):
    #     """Load configuration from pcdManip.json."""
    #     try:
    #         with open(self.config_path, 'r') as f:
    #             config = json.load(f)
            
    #         self.focused_bbox = config["focused_bbox"]
    #         self.bbox_min = np.array(self.focused_bbox[0])
    #         self.bbox_max = np.array(self.focused_bbox[1])
            
    #         logger.info(f"Loaded bounding box from {self.config_path}")
    #         logger.info(f"Bbox min: {self.bbox_min}")
    #         logger.info(f"Bbox max: {self.bbox_max}")
            
    #     except Exception as e:
    #         logger.error(f"Failed to load config from {self.config_path}: {e}")
    #         raise
            
    def is_point_in_bbox(self, points: np.ndarray) -> np.ndarray:
        """
        Check which points are inside the bounding box.
        
        Args:
            points: Point coordinates of shape (N, 3) or (N, 6) where first 3 are xyz
            
        Returns:
            Boolean mask of shape (N,) indicating which points are inside bbox
        """
        # Extract xyz coordinates (first 3 columns)
        xyz = points[:, :3]
        
        # Check if points are within bounding box
        inside_min = np.all(xyz >= self.bbox_min, axis=1)
        inside_max = np.all(xyz <= self.bbox_max, axis=1)
        inside_bbox = inside_min & inside_max
        
        return inside_bbox
        
    def trim_point_cloud_array(self, pcd_array: np.ndarray) -> np.ndarray:
        """
        Trim a single point cloud array based on bounding box.
        
        Args:
            pcd_array: Point cloud array of shape (N, 6) where 6 = [x, y, z, r, g, b]
            
        Returns:
            Trimmed point cloud array
        """
        if pcd_array.size == 0:
            return pcd_array
            
        # Get mask for points inside bounding box
        inside_mask = self.is_point_in_bbox(pcd_array)
        
        # Filter points
        trimmed_array = pcd_array[inside_mask]
        
        logger.debug(f"Trimmed from {pcd_array.shape[0]} to {trimmed_array.shape[0]} points")
        
        return trimmed_array


class TablePointCloudTrimmer:
    """
    Trims point cloud data based on table surface with holes for paper regions.
    """
    
    def __init__(self, table_center: np.ndarray, table_length: float, table_width: float, 
                 table_min_height: float, table_max_height: float, 
                 paper_bboxes: List[Tuple[np.ndarray, np.ndarray]]):
        """
        Initialize the table trimmer.
        
        Args:
            table_center: Center position of the table [x, y, z]
            table_length: Length of the table (x-direction)
            table_width: Width of the table (y-direction)
            table_min_height: Minimum height of the table surface
            table_max_height: Maximum height of the table surface
            paper_bboxes: List of (min, max) 2D bboxes for paper regions to exclude
        """
        self.table_center = np.array(table_center)
        self.table_length = table_length
        self.table_width = table_width
        self.table_min_height = table_min_height
        self.table_max_height = table_max_height
        self.paper_bboxes = paper_bboxes
        
        # Calculate table bounds
        half_length = table_length / 2
        half_width = table_width / 2
        
        self.table_x_min = table_center[0] - half_length
        self.table_x_max = table_center[0] + half_length
        self.table_y_min = table_center[1] - half_width
        self.table_y_max = table_center[1] + half_width
        
        logger.info(f"Table bounds: x[{self.table_x_min:.3f}, {self.table_x_max:.3f}], "
                   f"y[{self.table_y_min:.3f}, {self.table_y_max:.3f}], "
                   f"z[{self.table_min_height:.3f}, {self.table_max_height:.3f}]")
        logger.info(f"Paper holes: {len(paper_bboxes)} regions")
        
    def is_point_in_table_with_holes(self, points: np.ndarray) -> np.ndarray:
        """
        Check which points are inside the table surface but outside paper holes.
        
        Args:
            points: Point coordinates of shape (N, 3) or (N, 6) where first 3 are xyz
            
        Returns:
            Boolean mask of shape (N,) indicating which points are valid
        """
        # Extract xyz coordinates (first 3 columns)
        xyz = points[:, :3]
        
        # Check if points are within table bounds
        in_table_x = (xyz[:, 0] >= self.table_x_min) & (xyz[:, 0] <= self.table_x_max)
        in_table_y = (xyz[:, 1] >= self.table_y_min) & (xyz[:, 1] <= self.table_y_max)
        in_table_z = (xyz[:, 2] >= self.table_min_height) & (xyz[:, 2] <= self.table_max_height)
        in_table_bounds = in_table_x & in_table_y & in_table_z
        
        # Start with points that are within the table bounds
        table_mask = in_table_bounds.copy()

        # Remove paper regions from the table mask (we want to preserve papers)
        if len(self.paper_bboxes) > 0:
            for paper_min, paper_max in self.paper_bboxes:
                in_paper_x = (xyz[:, 0] >= paper_min[0]) & (xyz[:, 0] <= paper_max[0])
                in_paper_y = (xyz[:, 1] >= paper_min[1]) & (xyz[:, 1] <= paper_max[1])
                in_paper_region = in_paper_x & in_paper_y

                # Exclude paper regions from the table area
                table_mask = table_mask & ~in_paper_region

        # Valid points are those NOT in the remaining table mask
        valid_mask = ~table_mask

        return valid_mask
        
    def trim_point_cloud_array(self, pcd_array: np.ndarray) -> np.ndarray:
        """
        Trim a single point cloud array based on table surface with holes.
        
        Args:
            pcd_array: Point cloud array of shape (N, 6) where 6 = [x, y, z, r, g, b]
            
        Returns:
            Trimmed point cloud array
        """
        if pcd_array.size == 0:
            return pcd_array
            
        # Get mask for points in table but outside paper holes
        valid_mask = self.is_point_in_table_with_holes(pcd_array)
        
        # Filter points
        trimmed_array = pcd_array[valid_mask]
        
        logger.debug(f"Table trimmed from {pcd_array.shape[0]} to {trimmed_array.shape[0]} points")
        
        return trimmed_array


# Utility function for external use
def trimPcd(unifiedData: dict[str, np.ndarray], bbox_min: np.ndarray, bbox_max: np.ndarray, pcd_keys: list[str], param: str, *, surpress_print: bool = False) -> dict[str, np.ndarray]:
    """
    Trim point cloud data based on focused_bbox from pcdManip.json.
    
    Args:
        unifiedData: dict containing single frame data with keys:
                    ["rgb":np.ndarray,"depth":np.ndarray,"pointcloud":np.ndarray,
                     "segpointcloud":np.ndarray,"state":np.ndarray,"action":np.ndarray,
                     "endpose":np.ndarray,"cameras":dict,"gaussian":dict]
                    where pointcloud and segpointcloud are shape (Np, 6) for single frame
        bbox_min: Minimum bounding box coordinates (fallback if pcdManip file not found)
        bbox_max: Maximum bounding box coordinates (fallback if pcdManip file not found)
        pcd_keys: List of point cloud keys to trim
        param: Parameter to determine which pcdManip file to load (e.g., "1" for pcdManip1.json, None for pcdManip.json)
        surpress_print: If True, suppress console output (default: False)
        
    Returns:
        Dictionary containing original data with trimmed point cloud and segmented point cloud
    """
    if param is not None:
        pcd_manip_file = f'pcdManip{param}.json'
    else:
        pcd_manip_file = 'pcdManip.json'
    print(f"Loading pcdManip file: {pcd_manip_file}")
    try:
        pcd_manip_path = Path(__file__).parent / pcd_manip_file
        with open(pcd_manip_path, 'r') as f:
            pcd_manip_data = json.load(f)
        trim_bbox_min = np.array(pcd_manip_data["focused_bbox"][0])
        trim_bbox_max = np.array(pcd_manip_data["focused_bbox"][1])
        if not surpress_print:
            print(f"    📦 TRIM | Loaded bbox from {pcd_manip_file}")
    except FileNotFoundError:
        if not surpress_print:
            print(f"    ⚠️ TRIM | pcdManip file not found: {pcd_manip_file}, using fallback bbox")
        trim_bbox_min = bbox_min
        trim_bbox_max = bbox_max
 
    
    table_config, paper_bboxes = get_default_table_config(no_paper = True)
    BboxTrimmer = BboxPointCloudTrimmer(trim_bbox_min, trim_bbox_max)
    TableTrimmer = TablePointCloudTrimmer(table_config["table_center"], table_config["table_length"], table_config["table_width"], table_config["table_min_height"], table_config["table_max_height"], paper_bboxes)
    
    # Auto-trim all point cloud data where key contains "pointcloud"
    for pcd_key in pcd_keys:
        if pcd_key in unifiedData.keys():
            point_cloud_data = unifiedData[pcd_key]
            if point_cloud_data is not None:
                # Validate input shape
                if len(point_cloud_data.shape) != 2 or point_cloud_data.shape[1] != 6:
                    logger.warning(f"Skipping {pcd_key}: Expected shape (Np, 6), got {point_cloud_data.shape}")
                    return unifiedData
                original_count = point_cloud_data.shape[0]
                point_cloud_data = BboxTrimmer.trim_point_cloud_array(point_cloud_data)
                point_cloud_data = TableTrimmer.trim_point_cloud_array(point_cloud_data)
                unifiedData[pcd_key] = point_cloud_data

                trimmed_count = unifiedData[pcd_key].shape[0]
                
                # If trimmed to no points, add a point at (0,0,0)
                if trimmed_count == 0:
                    # Create a point at origin with default color (white: 255,255,255)
                    origin_point = np.array([[0.0, 0.0, 0.0, 255, 255, 255]], dtype=point_cloud_data.dtype)
                    unifiedData[pcd_key] = origin_point
                    trimmed_count = 1
                    if not surpress_print:
                        print(f"    ⚠️  Added origin point (0,0,0) to {pcd_key} as fallback")
                
                # --- Pretty console output -----------------------------------
                # Collect shapes for reporting
                if 'processed_shapes_in' not in locals():
                    processed_shapes_in = {}
                    processed_shapes_out = {}
                processed_shapes_in[pcd_key] = (original_count, 6) if len(point_cloud_data.shape)==2 else point_cloud_data.shape
                processed_shapes_out[pcd_key] = unifiedData[pcd_key].shape

    # After loop, print summary if any pointclouds processed
    if 'processed_shapes_in' in locals() and processed_shapes_in and not surpress_print:
        # Build inverse of pipeMap locally to find alias for single key case
        pipe_map_inv = {
            'pointcloud': 'p',
            'segpointcloud': 's',
            'repointcloud': 'r',
            'injpointcloud': 'i',
            'vggtpointcloud2': 'v2',
            'vggtpointcloud3': 'v3'
        }

        if len(processed_shapes_in) == 1:
            key = next(iter(processed_shapes_in))
            alias = pipe_map_inv.get(key, key)
            in_shape = processed_shapes_in[key]
            out_shape = processed_shapes_out[key]
            print(f"    ✂️  TRIM({alias}) | {in_shape} → {out_shape}")
        else:
            print(f"    ✂️  TRIM(*) | {processed_shapes_in} → {processed_shapes_out}")
 
    return unifiedData



# Default table configuration based on place_shoe.py and _base_task.py
def get_default_table_config(no_paper: bool = False):
    """
    Get default table configuration based on the place_shoe environment.
    
    Returns:
        Dictionary with table parameters and paper bboxes
    """
    # Table parameters from _base_task.py create_table_and_wall
    table_config = {
        "table_center": [0, 0, 0.74],  # Based on table_height=0.74
        "table_length": 1.21,
        "table_width": 0.71,
        "table_min_height": 0.74,  # Table surface height
        "table_max_height": 0.75,  # Slightly above table surface
    }
    
    if no_paper:
        # print("debug: no paper!")
        paper_bboxes = []
    else:
        # print("debug: with paper!")
    # Paper bbox from place_shoe.py target box
    # pose=[0, -0.08, 0.74], half_size=(0.13, 0.05, 0.0005)
        paper_center = [0, -0.08, 0.74]
        paper_half_size = [0.13, 0.05, 0.0005]
        
        paper_min = [paper_center[0] - paper_half_size[0], paper_center[1] - paper_half_size[1]]
        paper_max = [paper_center[0] + paper_half_size[0], paper_center[1] + paper_half_size[1]]
        
        paper_bboxes = [(np.array(paper_min), np.array(paper_max))]
    
    return table_config, paper_bboxes


